package main

import (
	"flag"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strconv"
	"syscall"
)

var (
	hostname string
)

func main() {
	flag.StringVar(&hostname, "hostname", "docker", "set hostname")
	flag.Parse()
	switch os.Args[1] {
	case "run":
		parent()
	case "child":
		child()
	default:
		panic("please input run")
	}
}

// parent 起一个进程运行命令，并且设置进程中的部分参数
func parent() {
	cmd := exec.Command("/proc/self/exe", append([]string{"child"}, os.Args[2:]...)...)
	// 读取/proc/self/exe可以获取当前程序的绝对路径
	cmd.SysProcAttr = &syscall.SysProcAttr{
		Cloneflags: syscall.CLONE_NEWUTS | syscall.CLONE_NEWPID | syscall.CLONE_NEWNS,
		// syscall.CLONE_NEWNET 需要实现网络桥接等机制
	}
	// UTS, PID and MNT namespaces 初始的挂载会继承
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	// 设置cgroups
	setupCgroups(cmd)

	run(cmd)
}

// child 实际运行命令的"容器"
func child() {
	cgroupPath := os.Getenv("CGROUP_CPU_PATH")
	if cgroupPath != "" {
		// 将当前进程加入到cgroup中
		selfPid := strconv.Itoa(os.Getpid())
		must(ioutil.WriteFile("/oldrootfs"+cgroupPath+"/tasks", []byte(selfPid), 0700))
	}

	must(syscall.Mount("rootfs", "rootfs", "", syscall.MS_BIND, ""))

	if err := syscall.Sethostname([]byte(hostname)); err != nil {
		fmt.Println("Failed to set hostname:", err)
		os.Exit(1)
	} //如果运行child，不加sudo会报错

	must(os.MkdirAll("rootfs/oldrootfs", 0700))
	must(syscall.PivotRoot("rootfs", "rootfs/oldrootfs"))
	must(os.Chdir("/")) //设置为/oldrootfs 为原来配置

	cmd := exec.Command(os.Args[2], os.Args[3:]...)
	cmd.Stdin = os.Stdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr

	run(cmd)
}

// setupCgroups 设置子进程资源限制
func setupCgroups(cmd *exec.Cmd) {
	// 创建cgroup
	cgroupPath := "/sys/fs/cgroup/cpu/container"
	os.MkdirAll(cgroupPath, 0755)

	// 将当前进程加入到cgroup中
	selfPid := strconv.Itoa(os.Getpid())
	must(ioutil.WriteFile(cgroupPath+"/tasks", []byte(selfPid), 0700))

	// 设置cgroup限制，这里示例设置CPU限制
	must(ioutil.WriteFile(cgroupPath+"/cpu.cfs_quota_us", []byte("10000"), 0700))

	// 将cgroup路径添加到cmd的环境变量中
	cmd.Env = append(cmd.Env, "CGROUP_CPU_PATH="+cgroupPath)
}

func must(err error) {
	if err != nil {
		panic(err)
	}
}

func run(cmd *exec.Cmd) {
	if err := cmd.Run(); err != nil {
		fmt.Println("ERROR", err)
		os.Exit(1)
	}
}
